import cv2
import numpy as np
import math
import rospy
from nav_msgs.msg import Odometry
import time
from collections import deque
from typing import Union, List


def get_RT(map):
    R, T = None, None
    if map == 'Town01':
        # Town01的变换矩阵(ok)
        R = np.array(
            [
                [1.608978933495218, -0.0006133637127350534],
                [0.002967129565265202, -1.616573483032402],
            ]
        )
        T = np.array([[21.96085734252182], [73.98491744019127]])
        err_xy = 5
        loacl_map_h_w = 140
    elif map == 'Town02':
        # Town02的变换矩阵(ok)
        R = np.array(
            [
                [2.853288729947334, -0.001954564367013564],
                [-0.01084418875167121, -2.902583603546258],
            ]
        )
        T = np.array([[75.19143327247443], [-271.2606785755725]])
        err_xy = 1.5
        loacl_map_h_w = 150
    elif map == 'Town03':
        # Town03的变换矩阵(ok very nice)
        R = np.array(
            [
                [1.42539846369879, -0.000613806695560062],
                [-0.006309288645396612, -1.42529645694144],
            ]
        )
        T = np.array([[266.8903605584413], [321.8950052342699]])
        # T = np.array([[265], [320]])
        err_xy = 5
        loacl_map_h_w = 150
    elif map == 'Town04':
        # Town04的变换矩阵(ok)
        R = np.array(
            [
                [0.6945775929658323, -0.00242031540532944],
                [-0.001151290440462326, -0.7332215587816061],
            ]
        )
        T = np.array([[374.2232704087302], [303.5616340310737]])
        err_xy = 5
        loacl_map_h_w = 150
    elif map == 'Town05':
        # Town05的变换矩阵(not ok 高架桥误差大)
        R = np.array(
            [
                [1.151724847113483, 0.04979182850230186],
                [-0.04002015696056613, -1.193283702860338],
            ]
        )
        T = np.array([[387.2353315551624], [325.3770351204567]])
        err_xy = 5
        loacl_map_h_w = 128
    elif map == 'Town06':
        # Town06的变换矩阵（ok）
        R = np.array(
            [
                [0.622355330577449, 0.005832891756840823],
                [-0.002292540060410386, -0.5624124260904821],
            ]
        )
        T = np.array([[242.799258978577], [278.7354311222649]])
        err_xy = 5
        loacl_map_h_w = 128
    elif map == 'Town07':
        # Town07的变换矩阵(not ok 误差大)
        R = np.array(
            [
                [1.641298667149452, -0.01306631310848138],
                [0.03929811782735146, -1.530002461260377],
            ]
        )
        T = np.array([[461.2455550308022], [399.1818917817909]])
        err_xy = 1
        loacl_map_h_w = 128
    return R, T


def transform_carla2img(loc, R, T):
    """carla世界坐标系 转 图像坐标系

    Args::
    Returns:
        img_loc (np shape=(2,1)): 图像坐标系
    """
    carla_loc = loc
    # img=R*carla+T
    img_loc = R.dot(carla_loc) + T
    u = math.floor(max(img_loc[0, 0], 0))
    v = math.floor(max(img_loc[1, 0], 0))
    img_loc = np.array([[u], [v]])
    return img_loc


def draw_map(map, loc):
    pix_x = loc[1][0]
    pix_y = loc[0][0]
    cv2.namedWindow("MAP", 0)
    cv2.resizeWindow("MAP", 800, 600)
    cv2.circle(map, (pix_x, pix_y), 5, (255, 0, 0), -1)  # BRG 蓝色-起点
    cv2.imshow("MAP", map)
    cv2.waitKey(0)


class RosNode:
    def __init__(self, node_name='navigation_node', debug_mode=False):
        self.node_name = node_name
        self.carla_odometry_topic = '/velodyne_points'
        self.is_odometry_received = False
        self.carla_odometry_msg = Odometry()
        self.msg_pool = {}

    def ros_init(self):
        rospy.init_node(self.node_name, anonymous=True, disable_signals=True)
        self.pub_waypoints = self.publisher()
        self.pool_regist(['odometry'])  # 注册消息池
        self.listener()

    def listener(self):
        rospy.Subscriber(
            self.carla_odometry_topic,
            Odometry,
            callback=self.callback_odometry,
            queue_size=1,
        )
        rospy.spin()

    def callback_odometry(self, data: Odometry):
        this_time = rospy.get_time() * 1000.0  # ms
        self.is_odometry_received = True
        self.carla_odometry_msg = data
        self.msg_pool_append("odometry", data, this_time)
        if self.debug_mode:
            print(f"get {type(data)} msg")

    def msg_pool_append(self, msg_name: str, msg_data, msg_time: float):
        """消息池存储函数

        Args:
            msg_name (str): _description_
            msg_data (_type_): 消息data
            msg_time (rostime): 消息对应的ros时间
        """
        msg_tuple = (msg_data, msg_time)
        try:
            self.msg_pool[msg_name].append(msg_tuple)
        except KeyError:
            self.msg_pool[msg_name].append(msg_tuple)

    def pool_regist(self, msg_name: Union[str, List[str]]):
        """消息池中字典注册

        Args:
            msg_name (Union[str, List[str]]): 消息名，或者包含消息名的列表
        """
        if isinstance(msg_name, list):
            for m in msg_name:
                self.pool_regist(m)
        elif isinstance(msg_name, str):
            self.msg_pool[msg_name] = deque(maxlen=5)


if __name__ == "__main__":
    # 1-离线demo
    R, T = get_RT(town='5')
    loc_carla = np.array([[99.63], [41.25]])
    loc_img = transform_carla2img(loc_carla, R, T)
    map = cv2.imread('/home/wrc/study/code/carla_tutorial/carlaMap/CARLAmap/Town05.jpg')
    draw_map(map, loc_img)

    # 2-ros主线程部分
    my_node = RosNode('carla_listen_node', debug_mode=False)
    my_node.ros_init()

    rate = rospy.Rate(10)
    while not rospy.is_shutdown():
        if my_node.is_odometry_received:
            carla_msg = my_node['odometry'][-1][0]
            carla_msg = Odometry()
            loc_x = carla_msg.pose.pose.position.x
            loc_y = carla_msg.pose.pose.position.y
            loc_carla = np.array([[loc_x], [loc_y]])
            loc_img = transform_carla2img(loc_carla, R, T)
            draw_map(map, loc_img)

        rate.sleep()
